Imports System.Text
Imports System.ComponentModel
Imports System.Windows.Forms
Imports System.Text.UnicodeEncoding
Imports System.Security.Cryptography
Imports System.IO

Public NotInheritable Class EncryptionWorker
  Inherits System.ComponentModel.BackgroundWorker

  Private bKey() As Byte
  Private bIV() As Byte
  Private bInitialised As Boolean = False
  Private rijM As RijndaelManaged = Nothing

  Private headerString As String = "CRYPTOR"
  Private headerBytes(7) As Byte
  Dim tempDir As String
  Dim tempFile As String

  Public Enum ReturnType As Integer
    Well = 0
    Badly = 1
    IncorrectPassword = 2
  End Enum

  Public Function GenerateHash(ByVal strSource As String) As String
    Return System.Convert.ToBase64String(New SHA384Managed().ComputeHash(New UnicodeEncoding().GetBytes(strSource)))
  End Function

  Public Sub Initialise(ByVal sPWH As String)
    'initialise rijM
    rijM = New RijndaelManaged
    'derive the key and IV using the 
    'PasswordDeriveBytes class
    Dim pdb As PasswordDeriveBytes = _
        New PasswordDeriveBytes(sPWH, _
        New MD5CryptoServiceProvider().ComputeHash(ConvertStringToBytes(sPWH)))
    'extract the key and IV
    bKey = pdb.GetBytes(32)
    bIV = pdb.GetBytes(16)

    'initialise headerBytes
    headerBytes = ConvertStringToBytes(headerString)

    With rijM
      .Key = bKey '256 bit key
      .IV = bIV '128 bit IV
      .BlockSize = 128 '128 bit BlockSize
      .Padding = PaddingMode.PKCS7
    End With

    bInitialised = True
    Me.WorkerReportsProgress = True
    Me.WorkerSupportsCancellation = True
  End Sub


  Private mResult As ReturnType
  Public Property Result() As ReturnType
    Get
      Return mResult
    End Get
    Set(ByVal value As ReturnType)
      mResult = value
    End Set
  End Property


  Private mInputFile As String
  Public Property InputFile() As String
    Get
      Return mInputFile
    End Get
    Set(ByVal value As String)
      mInputFile = value
    End Set
  End Property

  Private mOutputFile As String
  Public Property OutputFile() As String
    Get
      Return mOutputFile
    End Get
    Set(ByVal value As String)
      mOutputFile = value
    End Set
  End Property

  Private mEncrypt As Boolean
  Public Property Encrypt() As Boolean
    Get
      Return mEncrypt
    End Get
    Set(ByVal value As Boolean)
      mEncrypt = value
    End Set
  End Property
  Private mFilenameEncrypted As Boolean
  Public Property FilenameEncrypted() As Boolean
    Get
      Return mFilenameEncrypted
    End Get
    Set(ByVal value As Boolean)
      mFilenameEncrypted = value
    End Set
  End Property

  '  transform file
  Protected Overrides Sub OnDoWork(ByVal e As System.ComponentModel.DoWorkEventArgs)
    'make sure that all the initialisation has been completed:
    TransformFile(e)
  End Sub
  Public Overloads Sub TransformFile()
    Dim e As New DoWorkEventArgs(Nothing)
    TransformFile(e)
  End Sub
  ''' <summary>
  ''' set the filenames and the EncryptedFilename
  ''' </summary>
  ''' <param name="file"></param>
  ''' <remarks></remarks>
  Public Sub SetData(ByVal file As String, ByVal bEncrypt As Boolean, ByVal bFilenameEncrypted As Boolean)
    Dim path As String, filename As String
    path = IO.Path.GetDirectoryName(file)
    filename = IO.Path.GetFileName(file)
    Me.InputFile = file
    Me.FilenameEncrypted = bFilenameEncrypted
    Me.Encrypt = bEncrypt
    ' if filenames encrypted use a serial file number
    If bFilenameEncrypted And Me.Encrypt Then
      Dim i As Integer = 1, testFilename As String
      Do
        testFilename = IO.Path.Combine(path, String.Format("{0:000}.enc", i))
        i += 1
      Loop While IO.File.Exists(testFilename)
      Me.OutputFile = testFilename
    Else
      Me.OutputFile = tempFile
    End If
  End Sub
  Private Overloads Sub transformFile(ByVal e As System.ComponentModel.DoWorkEventArgs)
    If Not bInitialised Then
      Result = ReturnType.Badly : Return
    End If
    If Not IO.File.Exists(mInputFile) Then
      Result = ReturnType.Badly : Return
    End If

    Dim fsIn As FileStream = Nothing
    Dim fsOut As FileStream = Nothing
    Dim encStream As CryptoStream = Nothing
    Dim sw As IO.BinaryWriter = Nothing, sr As IO.BinaryReader = Nothing
    Dim retVal As ReturnType = ReturnType.Badly
    Dim sNameDecrypted As String = "encrypted.enc"
    Try
      'create the input and output streams:
      fsIn = New FileStream(mInputFile, FileMode.Open, FileAccess.Read)
      fsOut = New FileStream(mOutputFile, FileMode.Create, FileAccess.Write)

      'some helper variables
      Dim bBuffer(4096) As Byte '4KB buffer
      Dim bName() As Byte
      Dim lBytesRead As Long = 0
      Dim lFileSize As Long = fsIn.Length
      Dim lBytesToWrite, lNameLen As Integer

      ' ******************************** ENCRYPTION PERFORMED HERE ***************************
      If Encrypt Then
        encStream = New CryptoStream(fsOut, rijM.CreateEncryptor(bKey, bIV), CryptoStreamMode.Write)
        sw = New IO.BinaryWriter(encStream)
        'write the header to the output file for use when decrypting it
        sw.Write(headerBytes, 0, headerBytes.Length)
        If Me.FilenameEncrypted Then
          bName = Me.ConvertStringToBytes(mInputFile)
          sw.Write(bName.Length)
          encStream.Write(bName, 0, bName.Length)
        End If
        'this is the main encryption routine. it loops over the input data in blocks of 4KB,
        'and writes the encrypted data to disk
        Do
          If Me.CancellationPending Then
            e.Cancel = True
            Exit Try
          End If
          lBytesToWrite = fsIn.Read(bBuffer, 0, 4096)
          If lBytesToWrite = 0 Then Exit Do
          sw.Write(bBuffer, 0, lBytesToWrite)
          lBytesRead += lBytesToWrite
          Me.ReportProgress(CInt((lBytesRead / lFileSize) * 100))
        Loop
        ReportProgress(100)
        retVal = ReturnType.Well
        ' ******************************** DECRYPTION PERFORMED HERE ***************************
      Else
        encStream = New CryptoStream(fsIn, rijM.CreateDecryptor(bKey, bIV), CryptoStreamMode.Read)
        sr = New IO.BinaryReader(encStream)
        'read in the header
        Dim test(headerBytes.Length - 1) As Byte
        encStream.Read(test, 0, headerBytes.Length)
        Dim testStr As String

        'check to see if the file header reads correctly.
        'if it doesn't, then close the stream & jump out
        testStr = ConvertBytesToString(test)
        If testStr <> headerString Then
          encStream.Close()
          sr.Close()
          encStream.Clear()
          encStream = Nothing
          retVal = ReturnType.IncorrectPassword
          Exit Try
        End If
        If Me.FilenameEncrypted Then
          lNameLen = sr.ReadInt32
          sNameDecrypted = Me.ConvertBytesToString(sr.ReadBytes(lNameLen))
        End If

        'this is the main decryption routine. it loops over the input data in blocks of 4KB,
        'and writes the decrypted data to disk
        Do
          If Me.CancellationPending Then
            'if the cancel flag is set,
            'then jump out
            sr.Close() : encStream.Close()
            encStream.Clear()
            encStream = Nothing
            e.Cancel = True
            Exit Try
          End If
          lBytesToWrite = encStream.Read(bBuffer, 0, 4096)
          If lBytesToWrite = 0 Then Exit Do
          fsOut.Write(bBuffer, 0, lBytesToWrite)
          lBytesRead += lBytesToWrite
          Me.ReportProgress(CInt((lBytesRead / lFileSize) * 100))
        Loop
        Me.ReportProgress(100)
        retVal = ReturnType.Well
      End If
    Catch ex As Exception
    Finally
      'close all I/O streams (encStream first)
      If sr IsNot Nothing Then sr.Close()
      If sw IsNot Nothing Then sw.Close()
      If Not encStream Is Nothing Then
        encStream.Close()
      End If
      If Not fsOut Is Nothing Then
        fsOut.Close()
      End If
      If Not fsIn Is Nothing Then
        fsIn.Close()
      End If
    End Try
    'only delete the file if the password was bad, and
    'therefore its only an empty file
    Try
      If retVal = ReturnType.IncorrectPassword Then
        IO.File.Delete(mOutputFile)
      ElseIf retVal = ReturnType.Well Then ' delete input file, not needed
        If Me.FilenameEncrypted AndAlso Not Me.Encrypt Then
          Dim path As String = IO.Path.GetDirectoryName(mInputFile)
          Dim tempOut As String = IO.Path.Combine(path, sNameDecrypted)
          IO.File.Move(mOutputFile, tempOut)
          mOutputFile = tempOut
        End If
        IO.File.Delete(mInputFile)
      End If
    Catch ex As Exception
    End Try
  End Sub

  Public Function ConvertStringToBytes(ByVal sString As String) As Byte()
    Return New UnicodeEncoding().GetBytes(sString)
  End Function

  Public Function ConvertBytesToString(ByVal bytes() As Byte) As String
    Return New UnicodeEncoding().GetString(bytes)
  End Function

  Public Sub New()
    tempDir = Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData) & "\Cryptor"
    Dim di As New DirectoryInfo(tempDir)
    di.Create()
    tempFile = tempDir & "\~temp_encr"
  End Sub
End Class
